import asyncio
import logging
from typing import Tuple, Optional
from collections import deque
# 移除了错误的 PlaywrightException 导入，并从 playwright.async_api 导入 Browser
from playwright.async_api import async_playwright, Browser

from app.core.config import settings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Worker:
    """代表一个独立、可用的匿名会话工作单元。"""
    def __init__(self, cookie: str, nonce: str):
        self.cookie = cookie
        self.nonce = nonce
        self.failed_requests = 0

    def __repr__(self):
        return f"<Worker nonce='{self.nonce[:8]}...' failed={self.failed_requests}>"

class WorkerManager:
    """
    负责创建、维护和管理 Worker 池的单例类。
    """
    _instance = None
    _lock = asyncio.Lock()

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(WorkerManager, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, '_initialized'):
            self.pool: deque[Worker] = deque()
            self.browser: Optional[Browser] = None
            self._playwright = None
            self._maintenance_task: Optional[asyncio.Task] = None
            self._initialized = True
            self.total_created = 0
            self.total_failed = 0

    async def initialize(self):
        """启动 Playwright 并初始化浏览器。"""
        async with self._lock:
            if self.browser:
                return
            try:
                logger.info("正在启动 Playwright...")
                self._playwright = await async_playwright().start()
                self.browser = await self._playwright.chromium.launch(headless=True)
                logger.info("Playwright 浏览器实例已启动。")
                self._maintenance_task = asyncio.create_task(self._maintain_pool())
            except Exception as e:
                logger.critical(f"无法初始化 Playwright 或浏览器: {e}", exc_info=True)
                raise

    async def _create_worker(self) -> Optional[Worker]:
        """创建一个新的工作单元。"""
        if not self.browser:
            logger.error("浏览器未初始化，无法创建 Worker。")
            return None

        try:
            logger.info("正在创建新的 Worker...")
            context = await self.browser.new_context()
            page = await context.new_page()

            # 1. 访问目标页面以获取初始 Cookie
            await page.goto("https://writify.ai/tool/ai-chat/", timeout=settings.WORKER_CREATION_TIMEOUT * 1000)
            
            # 2. 调用 start_session API 获取 nonce
            api_url = "https://writify.ai/wp-json/mwai/v1/start_session"
            
            # 使用 page.evaluate 执行 fetch 请求，可以自动携带当前上下文的 cookies
            response_json = await page.evaluate(f"""
                fetch("{api_url}", {{ method: "POST" }})
                    .then(response => response.json())
            """)

            if response_json and response_json.get("success"):
                session_id = response_json.get("sessionId")
                rest_nonce = response_json.get("restNonce")
                
                # 从上下文中获取完整的 cookie 字符串
                all_cookies = await context.cookies()
                cookie_str = "; ".join([f"{c['name']}={c['value']}" for c in all_cookies])

                await context.close()

                if rest_nonce and cookie_str:
                    worker = Worker(cookie=cookie_str, nonce=rest_nonce)
                    logger.info(f"成功创建 Worker: {worker}")
                    self.total_created += 1
                    return worker
            
            await context.close()
            logger.error(f"创建 Worker 失败：未能从 start_session 获取有效数据。响应: {response_json}")
            self.total_failed += 1
            return None

        # 捕获更通用的异常以提高健壮性
        except Exception as e:
            logger.error(f"创建 Worker 时发生未知错误: {e}", exc_info=True)
            self.total_failed += 1
            # 确保在异常情况下也关闭 context
            if 'context' in locals() and not context.is_closed():
                await context.close()
            return None

    async def get_worker(self) -> Worker:
        """从池中获取一个可用的 Worker。"""
        async with self._lock:
            if not self.pool:
                # 如果池为空，尝试立即创建一个
                logger.warning("Worker 池为空，尝试紧急创建一个。")
                worker = await self._create_worker()
                if worker:
                    return worker
                raise RuntimeError("无法获取 Worker，池为空且无法创建新的。")
            
            # 轮询
            worker = self.pool.popleft()
            self.pool.append(worker)
            logger.info(f"提供 Worker: {worker}")
            return worker

    async def retire_worker(self, worker: Worker):
        """当一个 Worker 请求失败时，标记它并可能回收。"""
        async with self._lock:
            worker.failed_requests += 1
            logger.warning(f"Worker {worker} 失败次数增加: {worker.failed_requests}")
            # 简单策略：失败一次就回收
            if worker in self.pool:
                self.pool.remove(worker)
                logger.info(f"已从池中移除失败的 Worker: {worker}")

    async def _maintain_pool(self):
        """后台任务，持续维护 Worker 池的大小。"""
        while True:
            try:
                async with self._lock:
                    needed = settings.WORKER_POOL_SIZE - len(self.pool)
                
                if needed > 0:
                    logger.info(f"Worker 池需要补充 {needed} 个 Worker。")
                    for _ in range(needed):
                        worker = await self._create_worker()
                        if worker:
                            async with self._lock:
                                self.pool.append(worker)
                
                await asyncio.sleep(10) # 每10秒检查一次
            except asyncio.CancelledError:
                logger.info("维护任务被取消。")
                break
            except Exception as e:
                logger.error(f"维护任务出现异常: {e}", exc_info=True)
                await asyncio.sleep(30) # 发生异常后等待更长时间

    async def get_status(self) -> dict:
        """获取管理器状态。"""
        return {
            "pool_size": len(self.pool),
            "target_pool_size": settings.WORKER_POOL_SIZE,
            "workers": [repr(w) for w in self.pool],
            "total_created": self.total_created,
            "total_failed_creations": self.total_failed,
        }

    async def close(self):
        """关闭 Playwright 和浏览器。"""
        async with self._lock:
            if self._maintenance_task:
                self._maintenance_task.cancel()
            if self.browser:
                await self.browser.close()
            if self._playwright:
                await self._playwright.stop()
            logger.info("WorkerManager 已关闭。")

worker_manager = WorkerManager()
